// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/common/base/portability.h>
#include <pollux/dwio/common/column_visitors.h>
#include <pollux/dwio/common/direct_decoder.h>
#include <pollux/dwio/common/selective_column_reader.h>
#include <pollux/dwio/common/type_utils.h>
#include <pollux/exec/aggregation_hook.h>
#include <pollux/type/timestamp.h>
#include <pollux/vector/constant_vector.h>
#include <pollux/vector/dictionary_vector.h>
#include <pollux/vector/flat_vector.h>

#include <numeric>

namespace kumo::pollux::dwio::common {
    template<typename T>
    void SelectiveColumnReader::ensureValuesCapacity(vector_size_t numRows) {
        if (values_ && (isFlatMapValue_ || values_->unique()) &&
            values_->capacity() >=
            BaseVector::byteSize<T>(numRows) + simd::kPadding) {
            return;
        }
        values_ = AlignedBuffer::allocate<T>(
            numRows + (simd::kPadding / sizeof(T)), memoryPool_);
        rawValues_ = values_->asMutable<char>();
    }

    template<typename T>
    void SelectiveColumnReader::prepareRead(
        int64_t offset,
        const RowSet &rows,
        const uint64_t *incomingNulls) {
        const vector_size_t numRows = rows.back() + 1;
        readNulls(offset, numRows, incomingNulls);

        // We check for all nulls and no nulls. We expect both calls to
        // bits::isAllSet to fail early in the common case. We could do a
        // single traversal of null bits counting the bits and then compare
        // this to 0 and the total number of rows but this would end up
        // reading more in the mixed case and would not be better in the all
        // (non)-null case.
        allNull_ = nullsInReadRange_ &&
                   bits::isAllSet(
                       nullsInReadRange_->as<uint64_t>(), 0, numRows, bits::kNull);
        if (nullsInReadRange_ &&
            bits::isAllSet(
                nullsInReadRange_->as<uint64_t>(), 0, numRows, bits::kNotNull)) {
            nullsInReadRange_ = nullptr;
        }

        innerNonNullRows_.clear();
        outerNonNullRows_.clear();
        outputRows_.clear();
        // Is part of read() and after read returns getValues may be called.
        mayGetValues_ = true;
        numValues_ = 0;
        valueSize_ = sizeof(T);
        inputRows_ = rows;
        if (scanSpec_->filter() || hasDeletion()) {
            outputRows_.reserve(rows.size());
        }

        ensureValuesCapacity<T>(rows.size());
        if (scanSpec_->keepValues() && !scanSpec_->valueHook()) {
            valueRows_.clear();
            prepareNulls(rows, nullsInReadRange_ != nullptr);
        }
    }

    template<typename T, typename TVector>
    void SelectiveColumnReader::getFlatValues(
        const RowSet &rows,
        VectorPtr *result,
        const TypePtr &type,
        bool isFinal) {
        POLLUX_CHECK_NE(valueSize_, kNoValueSize);
        POLLUX_CHECK(mayGetValues_);
        if (isFinal) {
            mayGetValues_ = false;
        }

        if (allNull_) {
            if (isFlatMapValue_) {
                if (flatMapValueConstantNullValues_) {
                    flatMapValueConstantNullValues_->resize(rows.size());
                } else {
                    flatMapValueConstantNullValues_ =
                            std::make_shared<ConstantVector<TVector> >(
                                memoryPool_, rows.size(), true, type, T());
                }
                *result = flatMapValueConstantNullValues_;
            } else {
                *result = std::make_shared<ConstantVector<TVector> >(
                    memoryPool_, rows.size(), true, type, T());
            }
            return;
        }

        if (valueSize_ == sizeof(TVector)) {
            compactScalarValues<TVector, TVector>(rows, isFinal);
        } else if (sizeof(T) >= sizeof(TVector)) {
            compactScalarValues<T, TVector>(rows, isFinal);
        } else {
            upcastScalarValues<T, TVector>(rows);
        }
        valueSize_ = sizeof(TVector);
        if (isFlatMapValue_) {
            if (flatMapValueFlatValues_) {
                auto *flat = flatMapValueFlatValues_->as_unchecked<FlatVector<TVector> >();
                flat->unsafeSetSize(numValues_);
                flat->set_nulls(resultNulls());
                flat->unsafeSetValues(values_);
                flat->setStringBuffers(std::move(stringBuffers_));
            } else {
                flatMapValueFlatValues_ = std::make_shared<FlatVector<TVector> >(
                    memoryPool_,
                    type,
                    resultNulls(),
                    numValues_,
                    values_,
                    std::move(stringBuffers_));
            }
            *result = flatMapValueFlatValues_;
        } else {
            *result = std::make_shared<FlatVector<TVector> >(
                memoryPool_,
                type,
                resultNulls(),
                numValues_,
                values_,
                std::move(stringBuffers_));
        }
    }

    template<>
    void SelectiveColumnReader::getFlatValues<int8_t, bool>(
        const RowSet &rows,
        VectorPtr *result,
        const TypePtr &type,
        bool isFinal);

    template<typename T, typename TVector>
    void SelectiveColumnReader::upcastScalarValues(const RowSet &rows) {
        POLLUX_CHECK_LE(rows.size(), numValues_);
        POLLUX_CHECK(!rows.empty());
        if (!values_) {
            return;
        }
        POLLUX_CHECK_GT(sizeof(TVector), sizeof(T));
        // Since upcast is not going to be a common path, allocate buffer to copy
        // upcasted values to and then copy back to the values buffer.
        std::vector<TVector> buf;
        buf.resize(rows.size());
        T *typedSourceValues = reinterpret_cast<T *>(rawValues_);
        RowSet sourceRows;
        // The row numbers corresponding to elements in 'values_' are in
        // 'valueRows_' if values have been accessed before. Otherwise
        // they are in 'outputRows_' if these are non-empty (there is a
        // filter) and in 'inputRows_' otherwise.
        if (!valueRows_.empty()) {
            sourceRows = valueRows_;
        } else if (!outputRows_.empty()) {
            sourceRows = outputRows_;
        } else {
            sourceRows = inputRows_;
        }
        if (valueRows_.empty()) {
            valueRows_.resize(rows.size());
        }
        vector_size_t rowIndex = 0;
        auto nextRow = rows[rowIndex];
        auto *moveNullsFrom = shouldMoveNulls(rows);
        for (size_t i = 0; i < numValues_; i++) {
            if (sourceRows[i] < nextRow) {
                continue;
            }

            POLLUX_DCHECK(sourceRows[i] == nextRow);
            buf[rowIndex] = typedSourceValues[i];
            if (moveNullsFrom && rowIndex != i) {
                bits::setBit(rawResultNulls_, rowIndex, bits::isBitSet(moveNullsFrom, i));
            }
            valueRows_[rowIndex] = nextRow;
            rowIndex++;
            if (rowIndex >= rows.size()) {
                break;
            }
            nextRow = rows[rowIndex];
        }
        ensureValuesCapacity<TVector>(rows.size());
        std::memcpy(rawValues_, buf.data(), rows.size() * sizeof(TVector));
        numValues_ = rows.size();
        valueRows_.resize(numValues_);
        values_->setSize(numValues_ * sizeof(TVector));
    }

    template<typename T, typename TVector>
    void SelectiveColumnReader::compactScalarValues(
        const RowSet &rows,
        bool isFinal) {
        POLLUX_CHECK_LE(rows.size(), numValues_);
        POLLUX_CHECK(!rows.empty());
        if (!values_ || (rows.size() == numValues_ && sizeof(T) == sizeof(TVector))) {
            if (values_) {
                values_->setSize(numValues_ * sizeof(T));
            }
            return;
        }

        POLLUX_CHECK_LE(sizeof(TVector), sizeof(T));
        T *typedSourceValues = reinterpret_cast<T *>(rawValues_);
        TVector *typedDestValues = reinterpret_cast<TVector *>(rawValues_);
        RowSet sourceRows;
        // The row numbers corresponding to elements in 'values_' are in
        // 'valueRows_' if values have been accessed before. Otherwise
        // they are in 'outputRows_' if these are non-empty (there is a
        // filter) and in 'inputRows_' otherwise.
        if (!valueRows_.empty()) {
            sourceRows = valueRows_;
        } else if (!outputRows_.empty()) {
            sourceRows = outputRows_;
        } else {
            sourceRows = inputRows_;
        }
        if (valueRows_.empty()) {
            valueRows_.resize(rows.size());
        }

        vector_size_t rowIndex = 0;
        auto nextRow = rows[rowIndex];
        const auto *moveNullsFrom = shouldMoveNulls(rows);
        for (size_t i = 0; i < numValues_; ++i) {
            if (sourceRows[i] < nextRow) {
                continue;
            }

            POLLUX_DCHECK_EQ(sourceRows[i], nextRow);
            typedDestValues[rowIndex] = typedSourceValues[i];
            if (moveNullsFrom && rowIndex != i) {
                bits::setBit(rawResultNulls_, rowIndex, bits::isBitSet(moveNullsFrom, i));
            }
            if (!isFinal) {
                valueRows_[rowIndex] = nextRow;
            }
            ++rowIndex;
            if (rowIndex >= rows.size()) {
                break;
            }
            nextRow = rows[rowIndex];
        }

        numValues_ = rows.size();
        valueRows_.resize(numValues_);
        values_->setSize(numValues_ * sizeof(TVector));
    }

    template<>
    void SelectiveColumnReader::compactScalarValues<bool, bool>(
        const RowSet &rows,
        bool isFinal);

    inline int32_t sizeOfIntKind(TypeKind kind) {
        switch (kind) {
            case TypeKind::SMALLINT:
                return 2;
            case TypeKind::INTEGER:
                return 4;
            case TypeKind::BIGINT:
                return 8;
            default:
                POLLUX_FAIL("Not an integer TypeKind: {}", static_cast<int>(kind));
        }
    }

    template<typename T>
    void SelectiveColumnReader::filterNulls(
        const RowSet &rows,
        bool isNull,
        bool extractValues) {
        const bool isDense = rows.back() == rows.size() - 1;
        // We decide is (not) null based on 'nullsInReadRange_'. This may be
        // set due to nulls in enclosing structs even if the column itself
        // does not add nulls.
        auto *rawNulls =
                nullsInReadRange_ ? nullsInReadRange_->as<uint64_t>() : nullptr;
        if (isNull) {
            if (!rawNulls) {
                // The stripe has nulls but the current range does not. Nothing matches.
            } else if (isDense) {
                bits::forEachUnsetBit(
                    rawNulls, 0, rows.back() + 1, [&](vector_size_t row) {
                        addOutputRow(row);
                        if (extractValues) {
                            addNull<T>();
                        }
                    });
            } else {
                for (auto row: rows) {
                    if (bits::isBitNull(rawNulls, row)) {
                        addOutputRow(row);
                        if (extractValues) {
                            addNull<T>();
                        }
                    }
                }
            }
            return;
        }

        POLLUX_CHECK(
            !extractValues,
            "filterNulls for not null only applies to filter-only case");
        if (!rawNulls) {
            // All pass.
            for (auto row: rows) {
                addOutputRow(row);
            }
        } else if (isDense) {
            bits::forEachSetBit(rawNulls, 0, rows.back() + 1, [&](vector_size_t row) {
                addOutputRow(row);
            });
        } else {
            for (auto row: rows) {
                if (!bits::isBitNull(rawNulls, row)) {
                    addOutputRow(row);
                }
            }
        }
    }
} // namespace kumo::pollux::dwio::common
